#!/usr/bin/env python3
import os
import re
import json
import datetime
import clang.cindex
from clang.cindex import CursorKind, TypeKind, Config
from typing import List, Dict, Tuple, Optional

# 配置libclang路径（根据实际路径修改）
Config.set_library_file("/usr/lib/llvm-15/lib/libclang.so")  # 修改为你的实际路径

class CodeAnalyzer:
    def __init__(self, compile_args: List[str] = None):
        """
        初始化代码分析器
        
        参数:
            compile_args: 编译参数列表 (如 ['-I./include', '-DDEBUG'])
        """
        self.compile_args = compile_args or []
        try:
            self.index = clang.cindex.Index.create()
            self.struct_defs = {}  # {struct_name: {'definition': ..., 'file': ..., 'line': ...}}
        except Exception as e:
            raise RuntimeError(f"Failed to initialize Clang index: {str(e)}")

    def parse_file(self, file_path: str) -> Optional[clang.cindex.TranslationUnit]:
        """
        解析C/C++源文件
        
        参数:
            file_path: 源文件路径
            
        返回:
            成功返回TranslationUnit对象，失败返回None
        """
        try:
            return self.index.parse(file_path, args=self.compile_args)
        except Exception as e:
            print(f"Warning: Failed to parse {file_path}: {str(e)}")
            return None

    def extract_global_vars(self, tu: clang.cindex.TranslationUnit) -> List[Dict]:
        """
        提取全局变量
        
        参数:
            tu: TranslationUnit对象
            
        返回:
            全局变量信息列表，每个元素包含:
            - name: 变量名
            - type: 变量类型
            - location: 文件路径和行号
            - comment: 关联注释
        """
        globals = []
        if tu is None:
            return globals
            
        for node in tu.cursor.walk_preorder():
            try:
                if (node.kind == CursorKind.VAR_DECL and 
                    node.linkage == clang.cindex.LinkageKind.EXTERNAL):
                    var_info = {
                        'name': node.spelling,
                        'type': node.type.spelling,
                        'location': f"{node.location.file}:{node.location.line}",
                        'comment': self._get_associated_comment(node)
                    }
                    globals.append(var_info)
            except Exception as e:
                print(f"Warning: Error processing node in {tu.spelling}: {str(e)}")
                continue
        return globals

    def extract_heap_allocs(self, tu: clang.cindex.TranslationUnit) -> List[Dict]:
        """
        提取堆内存分配点
        
        参数:
            tu: TranslationUnit对象
            
        返回:
            堆分配信息列表，每个元素包含:
            - location: 文件路径和行号
            - call: 分配函数名 (如malloc/calloc/new)
            - type: 推断的类型
            - comment: 关联注释
            - source_code: 源代码片段
        """
        heap_allocs = []
        if tu is None:
            return heap_allocs
            
        for node in tu.cursor.walk_preorder():
            try:
                if node.kind == CursorKind.CALL_EXPR:
                    alloc_func = node.displayname.split('(')[0]
                    if alloc_func in ('malloc', 'calloc', 'realloc', 'new', 'operator new'):
                        heap_info = {
                            'location': f"{node.location.file}:{node.location.line}",
                            'call': alloc_func,
                            'type': self._infer_heap_type(node),
                            'comment': self._get_associated_comment(node),
                            'source_code': self._get_source_code(node)
                        }
                        heap_allocs.append(heap_info)
            except Exception as e:
                print(f"Warning: Error processing call in {tu.spelling}: {str(e)}")
                continue
        return heap_allocs

    def _infer_heap_type(self, call_node: clang.cindex.Cursor) -> str:
        """
        推断堆分配变量的类型
        
        参数:
            call_node: 分配函数调用节点
            
        返回:
            类型字符串，如 "int*" 或 "struct foo*"
        """
        try:
            # 1. 检查类型转换 (type*)malloc()
            for child in call_node.get_children():
                if child.kind == CursorKind.CSTYLE_CAST_EXPR:
                    return child.type.spelling
            
            # 2. 检查赋值操作左侧类型
            parent = call_node.semantic_parent
            if parent and parent.kind == CursorKind.BINARY_OPERATOR and parent.displayname == '=':
                for child in parent.get_children():
                    if child != call_node:  # 取赋值操作左侧
                        return child.type.spelling + '*'  # 添加指针修饰
                        
            # 3. 检查变量声明
            if parent and parent.kind == CursorKind.VAR_DECL:
                return parent.type.spelling
                
        except Exception as e:
            print(f"Warning: Error inferring heap type: {str(e)}")
            
        return "void*"  # 默认返回通用指针类型

    def extract_struct_defs(self, file_path: str):
        """
        提取结构体定义
        
        参数:
            file_path: 源文件路径
        """
        try:
            with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                content = f.read()
                
            # 匹配结构体定义（兼容typedef和非typedef形式）
            pattern = r'(/\*.*?\*/\s*)?(typedef\s+)?struct\s+(\w+)\s*\{([^}]*)\}(\s*\w*;)?'
            for match in re.finditer(pattern, content, re.DOTALL):
                comment = match.group(1) or ""
                struct_name = match.group(3)
                struct_body = match.group(4)
                
                self.struct_defs[struct_name] = {
                    'definition': struct_body.strip(),
                    'comment': comment.strip(),
                    'file': file_path,
                    'line': len(content[:match.start()].split('\n'))
                }
                
        except Exception as e:
            print(f"Warning: Error processing {file_path}: {str(e)}")

    def _get_associated_comment(self, node: clang.cindex.Cursor) -> str:
        """
        获取与节点关联的注释
        
        参数:
            node: AST节点
            
        返回:
            注释内容字符串
        """
        try:
            # 1. 直接关联的注释
            if node.raw_comment:
                return node.raw_comment.strip()
                
            # 2. 前驱注释
            tokens = list(node.get_tokens())
            if len(tokens) > 1:
                prev_token = tokens[-2]
                if prev_token.spelling.startswith(('/*', '//')):
                    return prev_token.spelling.strip()
                    
            # 3. 前一行注释
            if node.location.file:
                with open(node.location.file.name, 'r') as f:
                    lines = f.readlines()
                    line_num = node.location.line - 1
                    if line_num > 0 and lines[line_num-1].strip().startswith(('/*', '//')):
                        return lines[line_num-1].strip()
                        
        except Exception as e:
            print(f"Warning: Error getting comment: {str(e)}")
            
        return ""

    def _get_source_code(self, node: clang.cindex.Cursor) -> str:
        """
        获取节点对应的源代码片段
        
        参数:
            node: AST节点
            
        返回:
            源代码字符串
        """
        try:
            if node.location.file:
                with open(node.location.file.name, 'r') as f:
                    lines = f.readlines()
                    start = node.extent.start.line - 1
                    end = node.extent.end.line
                    return ''.join(lines[start:end]).strip()
        except:
            pass
        return ""

    def analyze_codebase(self, root_dir: str) -> Dict:
        """
        分析整个代码库
        
        参数:
            root_dir: 代码库根目录
            
        返回:
            分析结果字典，包含:
            - metadata: 元数据
            - global_variables: 全局变量列表
            - heap_allocations: 堆分配列表
            - struct_definitions: 结构体定义字典
        """
        results = {
            'metadata': {
                'source_directory': os.path.abspath(root_dir),
                'analysis_timestamp': str(datetime.datetime.now()),
                'compiler_args': self.compile_args
            },
            'global_variables': [],
            'heap_allocations': [],
            'struct_definitions': {}
        }
        
        # 1. 收集所有结构体定义
        print("[1/3] Collecting struct definitions...")
        for root, _, files in os.walk(root_dir):
            for file in files:
                if file.endswith(('.c', '.h', '.cpp', '.hpp')):
                    path = os.path.join(root, file)
                    self.extract_struct_defs(path)
                
            # 2. 分析源文件
        print("[2/3] Analyzing source files...")
        seen_globals = set()  # 用于记录已出现的全局变量标识符
        temp_global_vars = []  # 临时存储所有全局变量

        for root, _, files in os.walk(root_dir):
            for file in files:
                if file.endswith(('.c', '.cpp')):
                    path = os.path.join(root, file)
                    print(f"  Processing {path}...")
                    tu = self.parse_file(path)

                    # 提取全局变量（临时存储）
                    globals = self.extract_global_vars(tu)
                    for var in globals:
                        unique_key = (var['name'], var['type'], var['location'])
                        if unique_key not in seen_globals:
                            seen_globals.add(unique_key)
                            temp_global_vars.append(var)

        # 所有处理完成后，写入 results
        results['global_variables'].extend(temp_global_vars)

        # 提取堆分配
        for alloc in self.extract_heap_allocs(tu):
            # 关联结构体定义
            if 'struct' in alloc['type']:
                struct_name = re.search(r'struct\s+(\w+)', alloc['type'])
                if struct_name and struct_name.group(1) in self.struct_defs:
                    alloc['struct_definition'] = {
                        'name': struct_name.group(1),
                        'file': self.struct_defs[struct_name.group(1)]['file'],
                        'line': self.struct_defs[struct_name.group(1)]['line']
                    }
            results['heap_allocations'].append(alloc)
                
            # 3. 添加结构体定义
            print("[3/3] Finalizing results...")
            results['struct_definitions'] = self.struct_defs
            
            return results

    def save_to_json(self, data: Dict, output_file: str):
        """
        将分析结果保存为JSON文件
        
        参数:
            data: 分析结果字典
            output_file: 输出文件路径
        """
        try:
            with open(output_file, 'w') as f:
                json.dump(data, f, indent=2)
            print(f"Results saved to {output_file}")
        except Exception as e:
            print(f"Error saving results: {str(e)}")


if __name__ == "__main__":
    import argparse
    
    # 命令行参数解析
    parser = argparse.ArgumentParser(description='C/C++代码分析工具')
    parser.add_argument('source_dir',default='', help='源代码目录路径')
    parser.add_argument('-o', '--output', default='analysis_results.json', 
                       help='输出JSON文件路径 (默认: analysis_results.json)')
    parser.add_argument('-I', '--include', action='append', default=[], 
                       help='添加头文件搜索路径 (可多次使用)')
    parser.add_argument('-D', '--define', action='append', default=[], 
                       help='定义预处理宏 (可多次使用)')
    args = parser.parse_args()
    
    # 构建编译参数
    compile_args = []
    for inc in args.include:
        compile_args.append(f'-I{inc}')
    for define in args.define:
        compile_args.append(f'-D{define}')
    
    # 执行分析
    analyzer = CodeAnalyzer(compile_args=compile_args)
    results = analyzer.analyze_codebase(args.source_dir)
    analyzer.save_to_json(results, args.output)
    
    # 打印摘要
    print("\n=== 分析摘要 ===")
    print(f"全局变量: {len(results['global_variables'])}")
    print(f"堆内存分配: {len(results['heap_allocations'])}")
    print(f"结构体定义: {len(results['struct_definitions'])}")
    print(f"完整结果已保存到 {args.output}")